# -*- coding: utf-8 -*-

"""
@author: bw
@contact: bw@handsdriver.net
"""

import collections as C
import itertools   as I

import sys, os, fcntl, errno
import threading
import subprocess
import time
import random
import logging
import select
import hashlib

import plumbum.commands           as _plmcmd
import plumbum.commands.processes as _plumprocs
import plumbum.machines.base      as _plmbase
import plumbum.machines.session   as _plmsess


logger = logging.getLogger('plumbum.vagrant')

POPEN_TERMINATION_LIMIT = 10

STDIN_DEFAULT  = None
STDOUT_DEFAULT = subprocess.PIPE
STDERR_DEFAULT = subprocess.PIPE


class VagrantSession:
	inuse = 0
	busy  = property(lambda self: self.inuse != 0)

	current_cwd   = None
	current_env   = None
	current_popen = None

	endline = os.linesep.encode('ascii')

	ident_counter = I.count(1).__next__
	ident = None

	__owner = None

	def __init__(self, process, encoding='auto', connect_timeout=5):
		self.ident = self.ident_counter()

		self.process = process
#		self.custom_encoding = process.custom_encoding if encoding == 'auto' else encoding

		if connect_timeout:
			def closer():
				logger.error('Connection to %s timed out (%d sec)', proc, connect_timeout)
				self.close()

			timer = threading.Timer(connect_timeout, self.close)
			timer.start()

		try:
			self.run('')
		finally:
			if connect_timeout:
				timer.cancel()

		self.update_cwd()
		self.update_env()

	def __repr__(self):
		if self.ident is None:
			return '<VagrantSession:NOIDENT>'

		try:
			alive = self.alive()
			return '<VagrantSession:{} inuse={} {}{}>'.format(self.ident, self.inuse,
				'alive' if alive else 'DEAD',
				' BROKEN' if self.broken() else '')
		except Exception:
			logger.exception('Error in repr of VagrantSession')
			return '<VagrantSession:FAILURE>'

	def __enter__(self):
		if self.broken():
			raise _plmsess.ShellSessionError('I\'m broken, possible stdin closed')
		assert self.acquire(), (self.inuse, self.current_popen)
		return self

	def __exit__(self, etype, evalue, tb):
		self.release()
		if etype is not None or self.broken():
			self.close()

	def broken(self):
		if self.process is None:
			return True
		return self.process.stdin.closed \
		    or self.process.stdout.closed\
		    or self.process.stderr.closed

	def alive(self):
		# процесс в процессе, удалённый
		if self.current_popen and self.current_popen.poll() is None:
			return True
		# или локальный (vagrant ssh)
		return self.process and self.process.poll() is None

	def close(self):
		logger.debug('Close %r', self)

		if self.alive():
			if self.current_popen and self.current_popen.poll() is None:
				try:
					self.current_popen.close(timeout=5.0)
				except Exception:
					logger.exception('Unhandled exception by close %r', self.current_popen)

			if self.process is not None:
				if not self.process.stdin.closed:
					try:
						self.process.stdin.write(b'\nexit\n\n\nexit\n\n'.replace(b'\n', self.endline))
						self.process.stdin.flush()
					except Exception:
						logger.exception('Unhandled exception during exit %r', self.process)

				time.sleep(0.05)

				for pipe in [self.process.stdin, self.process.stdout, self.process.stderr]:
					try:
						pipe.close()
					except Exception:
						logger.exception('Unhandled exception during close %r of %r', pipe, self.process)

				try:
					self.process.terminate()
				except Exception:
					logger.exception('Unhandled exception during terminate %r', self.process)

			for counter in I.count(-5):
				if not self.alive():
					break

				if 0 <= counter:
					sys.stderr.write('.')
					sys.stderr.flush()

				if 100 <= counter:
					sys.stderr.write('FORCED EXIT\n')
					sys.stderr.flush()
					os._exit(99)

				time.sleep(0.1 if counter < 50 else 0.5)

		self.process = None

	def acquire(self):
		if self.broken():
			return False

		if self.current_popen is None or self.current_popen.returncode is not None:
			thread_ident = threading.get_ident()

			if self.inuse == 0:
				self.inuse += 1
				self.__owner = thread_ident
				return True

			elif self.__owner == thread_ident:
				self.inuse += 1
				return True

		return False

	def release(self):
		if 0 < self.inuse:
			thread_ident = threading.get_ident()
			assert self.__owner == thread_ident, (self.__owner, thread_ident)
			self.inuse -= 1
			if self.inuse == 0:
				del self.__owner

	def update_cwd(self):
		self.current_cwd = self.run('pwd')[1].strip()
		return self.current_cwd

	def update_env(self):
		env0 = self.run('env -0; echo')

		if env0[0] == 0 and not env0[2].rstrip():
			self.current_env = dict(line.split('=', 1) for line in env0[1].split('\x00') if '=' in line)

		else:
			lines = self.run('env; echo')[1].splitlines()
			split = (line.split('=', 1) for line in lines)
			keys  = (line[0] for line in split if 1 < len(line))
			runs  = ((key, self.run('printenv "' + key + '"; echo')) for key in keys)

			self.current_env = {key: run[1].rstrip('\n') for (key, run) in runs
				if run[0] == 0 and run[1].rstrip('\n') and not run[2]}

		return self.current_env

	def popen(self, args, cbdone=None, stdin=STDIN_DEFAULT, stdout=STDOUT_DEFAULT, stderr=STDERR_DEFAULT):
		if self.broken():
			raise _plmsess.ShellSessionError('I\'m broken, possible stdin closed')

		if self.process is None:
			raise _plmsess.ShellSessionError('Shell session has already been closed')

		if self.current_popen and self.current_popen.returncode is None:
			raise _plmsess.ShellSessionError('Each shell may start only one process at a time')

		def ondone(exc):
			# @xxx: может быть вызвано дважды, см. L{VagrantPopen.__stdio_bridge_loop}
			self.inuse -= 1
			if cbdone:
				cbdone(exc)

		self.inuse += 1

		try:
			self.current_popen = VagrantPopen(self, args, cbdone=ondone, stdin=stdin, stdout=stdout, stderr=stderr)
			logger.debug('Running %r', self.current_popen)
		except BaseException as exc:
			ondone(exc)
			raise

		return self.current_popen

	def run(self, cmd, retcode=0):
		with self:
			proc = self.popen(cmd)
			try:
				return _plumprocs.run_proc(proc, retcode)
			finally:
				for fobj in [proc.stdin, proc.stdout, proc.stderr]:
					if fobj is not None:  # not PIPE
						try:
							fobj.close()
						except OSError:  # @xxx: ...
							logger.exception('Unhandler exception during close %r of %r', fobj, proc)


class VagrantPopen(_plmbase.PopenAddons):
	__done      = False
	__stdout_wr = None
	__stderr_wr = None
	__marker    = None

	ident_counter = I.count(1).__next__

	ident   = None
	cmdline = None

	failure    = None
	returncode = None
	stdin  = None
	stdout = None
	stderr = None

	custom_encoding = sys.getdefaultencoding()
	endline = os.linesep.encode('ascii')

	# @xxx: plumbum.commands.base:Pipeline.popen
	_proc = property(lambda self: self)

	def __init__(self, session, args, cbdone=None, stdin=STDIN_DEFAULT, stdout=STDOUT_DEFAULT, stderr=STDERR_DEFAULT):
		# @xxx: нужно для `run_proc` -> `PopenAddons` -> `ProcessExecutionError`
		# @xxx: должен быть list/tuple, но см. `ShellSession.popen`

		self.session = session
		self.cbdone  = cbdone

		if isinstance(args, C.Sequence) and not isinstance(args, (bytes, str)):
			self.argv = list(args)
			cmdline = ' '.join(_plmcmd.shquote_list(self.argv))  # ?????????????
		elif isinstance(args, bytes):
			cmdline = args.decode(self.custom_encoding)
			self.argv = [cmdline]
		else:
			cmdline = str(args)
			self.argv = [cmdline]

		self.ident = self.ident_counter()

		marker  = 'stop-marker:{}:'.format(self.ident)
#		marker += hashlib.sha1('{}.{}'.format(time.time(), random.random()).encode('ascii')).hexdigest()[:10]
		marker += str(random.randint(10**9, 10**10-1))

		self.__marker = marker.encode('ascii')

		if not cmdline.strip():
			cmdline = 'true'

		self.cmdline = '{cmdline}; echo "[[{marker},exitcode=$?]]"; echo "[[{marker}]]" 1>&2'\
			.format(cmdline=cmdline, marker=marker)

		cmdline = self.cmdline.encode(self.custom_encoding)

		assert stdin is None or hasattr(stdin, 'fileno'), (self, stdin)

		session.process.stdin.write(cmdline + b'\n')
		session.process.stdin.flush()

		if stdin is None:
			self.stdin = os.fdopen(os.dup(session.process.stdin.fileno()))
		else:
			self.stdin = os.fdopen(os.dup(stdin.fileno()))

		self.__session_stdin  = session.process.stdin
		self.__session_stdout = session.process.stdout
		self.__session_stderr = session.process.stderr

		self.__kick_rfd, self.__kick_wfd = os.pipe()

		if stdout is None:
			stdout_rfd = None
			stdout_wfd = os.dup(sys.stdout.fileno())
		elif stdout is subprocess.PIPE:
			stdout_rfd, stdout_wfd = os.pipe()
		elif hasattr(stdout, 'write'):
			stdout_rfd = None
			stdout_wfd = os.dup(stdout.fileno())
		else:
			raise ValueError('Unknown `stdout`', stdout)

		if stderr is None:
			stderr_rfd = None
			stderr_wfd = os.dup(sys.stderr.fileno())
		elif stderr is subprocess.PIPE:
			stderr_rfd, stderr_wfd = os.pipe()
		elif hasattr(stderr, 'write'):
			stderr_rfd = None
			stderr_wfd = os.dup(stderr.fileno())
		else:
			raise ValueError('Unknown `stderr`', stderr)

#		fcntl.fcntl(self.__kick_rfd, fcntl.F_SETFL, fcntl.fcntl(self.__kick_rfd, fcntl.F_GETFL) | os.O_NONBLOCK)
#		fcntl.fcntl(stdout_wfd, fcntl.F_SETFL, fcntl.fcntl(stderr_wfd, fcntl.F_GETFL) | os.O_NONBLOCK)
#		fcntl.fcntl(stderr_wfd, fcntl.F_SETFL, fcntl.fcntl(stderr_wfd, fcntl.F_GETFL) | os.O_NONBLOCK)

		self.__stdout_wr = os.fdopen(stdout_wfd, 'wb', buffering=0)
		self.__stderr_wr = os.fdopen(stderr_wfd, 'wb', buffering=0)

		self.stdout = os.fdopen(stdout_rfd, 'rb', buffering=0) if stdout_rfd is not None else None
		self.stderr = os.fdopen(stderr_rfd, 'rb', buffering=0) if stderr_rfd is not None else None

		# @todo: (?) weakref
		self.__stdio_thread = threading.Thread(
			name  ='VagrantPopen/stdio_bridge_loop',
			target=self.__stdio_bridge_loop)
		self.__stdio_thread.start()

#		except:
#			os.close(stdout_rfd)
#			os.close(stdout_wfd)
#			os.close(stderr_rfd)
#			os.close(stderr_wfd)
#			raise

	def __repr__(self):
		if self.ident is None or self.session is None or self.session.ident is None:
			return '<VagrantPopen:NOIDENT>'
		return '<VagrantPopen:{}.{} {}>'.format(self.session.ident, self.ident,
			repr(self.cmdline) if self.cmdline is not None else 'NOCMDLINE')

	def __stdio_bridge_loop(self):
		# @xxx: это всё порнуха, но пока хоть как-нибудь, например, все дескрипторы блокируемые,
		#       а так же читается строка целиком, из-за `marker`-а,
		#       при этом возможна блокировка/зависание: сессионные stdout/stderr не закрываются при "завершении" работы

		evalue = None

		try:
			session_stdin  = self.__session_stdin
			session_stdout = self.__session_stdout
			session_stderr = self.__session_stderr

			stdin_fileno = self.stdin.fileno()

			marker_prefix = b'[[' + self.__marker

			read_streams = {
				self.__kick_rfd        : None,
				stdin_fileno           : (self.stdin,     session_stdin,  None),
				session_stdout.fileno(): (session_stdout, self.__stdout_wr, []),
				session_stderr.fileno(): (session_stderr, self.__stderr_wr, []),
			}

			termination  = False
			term_timeout = None

			while not termination:
#				logger.debug('VagrantPopen.select > (kick=%d) %r', self.__kick_rfd, sorted(read_streams.keys()))
				rfds, wfsd, xfds = select.select(list(read_streams.keys()), [], [])
#				logger.debug('VagrantPopen.select < %r %r %r', sorted(rfds), sorted(wfsd), sorted(xfds))

				for rfd in rfds:
					if rfd == self.__kick_rfd:
						if self.returncode is not None or term_timeout is not None and term_timeout <= time.time():
							termination = True
							os.read(rfd, 1)
						elif term_timeout is None:
							term_timeout = time.time() + POPEN_TERMINATION_LIMIT
					else:
						pipein, pipeout, buffer = read_streams[rfd]
						chunk = os.read(rfd, 256)  # @todo: [OSError] (9) Bad file descriptor

#						logger.debug('VagrantPopen.pipein:%d %r', rfd, chunk)

						if buffer is None:
							if chunk:
								pipeout.write(chunk)
							else:
								del read_streams[rfd]
								pipeout.close()
							continue

						lines = []

						if not chunk:
							del read_streams[rfd]
							lines.append(b''.join(buffer))

						elif self.endline in chunk:
							if buffer:
								chunk = b''.join(buffer) + chunk
								del buffer[:]

							while self.endline in chunk:
								line, chunk = chunk.split(self.endline, 1)
								lines.append(line)

							if chunk:
								buffer[:] = [chunk]

						else:
							buffer.append(chunk)

						for no, line in enumerate(lines):
							if marker_prefix in line:
								line, tail = line.split(marker_prefix, 1)
								# tail = ']]...' or ',exitcode=XXX]]...'

								if tail[:10] == b',exitcode=':
									exitcode, tail = tail[10:].split(b']]', 1)
									exitcode = int(exitcode)

									if self.returncode is None:
										self.returncode = exitcode
									else:
										logger.error('`returncode` already exist (%r), ignore another (%r)',
											self.returncode, exitcode)

								elif tail[:2] == b']]':
									tail = tail[2:]

								else:
									logger.error('Unexpected marker terminator: %r', tail)

								if tail:
									logger.error('Lost: %r', tail)

								if line:
									logger.error('Unexpected line: %r', line)
									if not pipeout.closed:
										pipeout.write(line + self.endline)  # @xxx: возможна блокировка

								read_streams.pop(rfd, None)

							elif not pipeout.closed:
								pipeout.write(line + self.endline)  # @xxx: возможна блокировка
								# @xxx: (32) Broken pipe

							else:
								logger.error('Lost: %r', line)

						if rfd not in read_streams:
							pipeout.close()

				if len(read_streams) == 1 and (self.__kick_rfd in read_streams or stdin_fileno in read_streams):
					termination = True
				elif len(read_streams) == 2 and self.__kick_rfd in read_streams and stdin_fileno in read_streams:
					termination = True

		except BaseException as exc:
			self.failure = exc
			logger.exception('Unhandled exception in `VagrantPopen.__stdio_bridge_loop`')

		finally:
			try:
				try:
					os.close(self.__kick_rfd)
				except OSError as exc:
					if exc.errno != errno.EBADF:
						raise

				self.__stdout_wr.close()
				self.__stderr_wr.close()
			finally:
				self.__done = True

				if self.returncode is None:
					self.returncode = 99

				# не хочу делать это в потоке
				#if self.cbdone is not None:
				#	self.cbdone(evalue)
				#	self.cbdone = None

	iter_lines = _plumprocs.iter_lines

	def __iter__(self):
		return self.iter_lines()

	def close(self, timeout=None):
		logger.debug('Close %r', self)

		try:
			try:
				os.write(self.__kick_wfd, b'x')
			except OSError as exc:
				if exc.errno not in (errno.EBADF, errno.EPIPE):
					raise

			try:
				os.close(self.__kick_wfd)
			except OSError as exc:
				if exc.errno != errno.EBADF:
					raise

			if self.__stdio_thread.isAlive():
				if timeout is not None:
					stoptime = time.time() + max(0.001, timeout)

				time.sleep(0.01)

				if self.__stdio_thread.isAlive():
					logging.debug('Waiting stdio thread...')

					while self.__stdio_thread.isAlive():
						time.sleep(0.01)

						if timeout is not None and stoptime <= time.time():
							raise TimeoutError()
		finally:
			self.__stdout_wr.close()
			self.__stderr_wr.close()

			self.stdin.close()

			if self.stdout is not None:
				self.stdout.close()
				self.stdout = None

			if self.stderr is not None:
				self.stderr.close()
				self.stderr = None

			self.__done = True

			if self.__session_stdin.closed:
				pass

			if self.cbdone is not None:
				try:
					self.cbdone(self.failure)
				except Exception:
					logger.exception('Unhandled exception in Popen callback')
				self.cbdone = None

	def kill(self):
		try:
			self.close(timeout=2.0)
		except TimeoutError:
			self.session.process.kill()
			self.__done = True

	def terminate(self):
		try:
			self.close(timeout=2.0)
		except TimeoutError:
			self.session.process.terminate()
			self.__done = True

	def poll(self):
		"""Returns the process' exit code or ``None`` if it's still running"""
		if self.__done:
			return self.returncode

	def wait(self, timeout=None):
		"""Waits for the process to terminate and returns its exit code"""
		if not self.__done:
			self.communicate(timeout=timeout)
		return self.returncode

	def communicate(self, input=None, timeout=None):
		assert timeout is None

		sources = []

		if self.stdout is not None:
			# `stdout` is PIPE
			stdout_lines = []
			sources.append(('stdout', stdout_lines, self.stdout))
		else:
			stdout_lines = None

		if self.stderr is not None:
			# `stderr` is PIPE
			stderr_lines = []
			sources.append(('stderr', stderr_lines, self.stderr))
		else:
			stderr_lines = None

		try:
			i = 0

			while sources or input:
				if input:
					chunk = input[:256]
					input = input[256:]
					self.__session_stdin.write(chunk)
					self.__session_stdin.flush()

				i = (i + 1) % len(sources)
				name, lines, pipe = sources[i]

				try:
					line = pipe.readline()
					logger.debug('VagrantPopen.communicate: %s> %r', name, line)
				except EOFError:
					logger.debug('VagrantPopen.communicate: %s> Nothing returned', name)

					self.session.process.poll()
					returncode = self.session.process.returncode

					if returncode == 5:
						raise _plmsess.IncorrectLogin('Incorrect username or password provided')

					if returncode == 6:
						raise _plmsess.HostPublicKeyUnknown('The authenticity of the host can\'t be established')

					if name == 'stderr':
						raise _plmsess.SSHCommsChannel2Error('No stderr result detected. Does the remote have Bash as the default shell?')

					raise _plmsess.SSHCommsError('No communication channel detected. Does the remote exist?')

				if not line:
					del sources[i]
				else:
					lines.append(line)

		finally:
			self.close(timeout=10.0)
			if self.returncode is None:
				raise RuntimeError('Undefined `returncode`')

		return\
			b''.join(stdout_lines) if stdout_lines is not None else None,\
			b''.join(stderr_lines) if stderr_lines is not None else None
